#include "pch.h"
#include "Tracker.h"
#include <math.h>

Tracker::Tracker() {
	// pattern = cv::Mat(CORE_SIZE, CORE_SIZE, CV_64FC3, buffer);
	patterns[0] = cv::Mat(CORE_SIZE, CORE_SIZE, CV_32FC1, cv::Scalar(1.0)); // H
	patterns[1] = cv::Mat(CORE_SIZE, CORE_SIZE, CV_32FC1, cv::Scalar(0.0)); // S
	patterns[2] = cv::Mat(CORE_SIZE, CORE_SIZE, CV_32FC1, cv::Scalar(0.0)); // V
	last_hit = cv::Rect(0, 0, 0, 0);
}

#define H_MIN_1		0
#define H_MAX_1		10
#define H_MIN_2		156
#define H_MAX_2		180
#define S_MIN		43
#define S_MAX		255
#define V_MIN		46
#define V_MAX		255
#define SCORE_MIN   1000
double red_score(int h, int s, int v) {
	bool is_red = true;
	is_red = is_red && ((H_MIN_1 < h) && (h < H_MAX_1) || (H_MIN_2 < h) && (h < H_MAX_2));
	is_red = is_red && ((S_MIN < s) && (s < S_MAX));
	is_red = is_red && ((V_MIN < v) && (v < V_MAX));
	if (is_red) {
		return 1.0;
	}
	else {
		return 0.0;
	}
}

double mat_red_score(cv::Mat target) {
	cv::Mat hsv[3];
	cv::split(target, hsv);
	cv::Mat h = hsv[0], s = hsv[1], v = hsv[2];
	cv::Mat is_red_1, is_red_2, is_red;
	cv::bitwise_and(H_MIN_1 < h, h < H_MAX_1, is_red_1);
	cv::bitwise_and(H_MIN_2 < h, h < H_MAX_2, is_red_2);
	cv::bitwise_or(is_red_1, is_red_2, is_red);
	cv::bitwise_and(S_MIN < s, s < S_MAX, is_red_1);
	cv::bitwise_and(is_red_1, is_red, is_red);
	cv::bitwise_and(V_MIN < s, s < V_MAX, is_red_1);
	cv::bitwise_and(is_red_1, is_red, is_red);
	is_red = is_red / 255;
	double score = cv::sum(is_red).val[0]; // 1 channel here
	return score;
}

// TODO: add location feature for score calculation
double Tracker::get_score(cv::Mat target, int channel) {
	return mat_red_score(target);

	double score = 0;
	for (int i = 0; i < target.rows; ++i) {
		for (int j = 0; j < target.cols; ++j) {
			cv::Vec3b hsv = target.at<cv::Vec3b>(i, j);
			int h = hsv.val[0];
			int s = hsv.val[1];
			int v = hsv.val[2];
			score += red_score(h, s, v);
		}
	}
	if (score > SCORE_MIN)
		return score;
	else
		return -1;
}

bool Tracker::valid_last_hit() {
	return (last_hit.width != 0) && (last_hit.height != 0);
}

// frame should be of RGB format
int Tracker::tracker_update(cv::Mat _frame) {
	cv::Mat frame;
	cv::cvtColor(_frame, frame, cv::COLOR_BGR2HSV);
	int width = frame.size().width;
	int height = frame.size().height;
	// cv::Mat channels[CHANNEL_SIZE];
	// cv::split(frame, channels);
	// for (int c = 0; c < CHANNEL_SIZE; ++c) {
	//     channels[c].convertTo(channels[c], CV_32FC1);
	// }
	int chunk_width = CORE_SIZE;
	int chunk_height = CORE_SIZE;
	int step_width, step_height;
	int width_start, height_start, width_end, height_end;
	double max_score = 0;

	if (valid_last_hit()) {
		step_width = FINE_STEP, step_height = FINE_STEP;
		int center_x = last_hit.x + last_hit.width / 2;
		int center_y = last_hit.y + last_hit.height / 2;
		width_start = center_x - SEARCH_SIZE / 2;
		if (width_start < 0) width_start = 0;
		height_start = center_y - SEARCH_SIZE / 2;
		if (height_start < 0) height_start = 0;
		width_end = width_start + SEARCH_SIZE;
		while (width_end > width-chunk_width) width_end = width_end - step_width;
		height_end = height_start + SEARCH_SIZE;
		while (height_end > height-chunk_width) height_end = height_end - step_height;
	}
	else {
		step_width = CORE_SIZE, step_height = CORE_SIZE;
		width_start = (width % chunk_width) / 2;
		height_start = (height % chunk_height) / 2;
		width_end = width - width_start;
		height_end = height - height_start;
	}

	cv::Rect max_rect(0, 0, 0, 0);
	// std::cout << width_start << width_end << height_start << height_end << std::endl;
	for (int i = height_start; i < height_end; i += step_height) {
		for (int j = width_start; j < width_end; j += step_width) {
			cv::Rect roi(j, i, chunk_width, chunk_height);
			cv::Mat target = frame(roi);
			double score = get_score(target, CHANNEL_SIZE);
			// std::cout << score << std::endl;
			if (score > max_score) {
				max_score = score;
				max_rect = roi;
			}
		}
	}
	last_score = max_score;
	last_hit = max_rect;
	return 0;
}

int Tracker::get_target(cv::Rect &roi, double &score) {
	if (!valid_last_hit()) {
		return -1;
	}
	roi = last_hit;
	score = last_score;
	return 0;
}